#ifndef EMB_RK_HPP
#define EMB_RK_HPP

#include <cmath>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <vector>

#include "IVPSolver.hpp"

class EmbRK : public IVPSolver {
public:
  EmbRK(int dim, double dt, double t0, double t1)
      : IVPSolver(dim, dt, t0, t1) {
    // 初始化时只分配一个位置用于初始条件
    m_u.resize(1, m_dim);
  }

  void solve(const Eigen::VectorXd &u0,
             const std::function<Eigen::VectorXd(Eigen::VectorXd, double)> &f) {
    if (u0.size() != m_dim) {
      std::cerr << "Invalid input size." << std::endl;
      return;
    }

    // 储存结果的向量
    std::vector<Eigen::VectorXd> solution;
    std::vector<double> timePoints;
    
    // 储存初始条件
    m_u.row(0) = u0;
    solution.push_back(u0);
    timePoints.push_back(m_t0);

    Eigen::MatrixXd Y(m_s, m_dim);
    
    double step = m_dt;
    double t = m_t0;
    
    // 设置误差容忍度
    double Eabs = 1e-6;  // 绝对误差容忍度
    double Erel = 1e-6;  // 相对误差容忍度
    
    // 步长控制参数
    const double rhoMax = 2.0;    // 最大步长增长因子
    const double rhoMin = 0.2;    // 最小步长缩减因子
    const double rho = 0.8;       // 安全因子
    const double q = std::min(m_pl, m_ph);  // 阶数
    
    // 当前解
    Eigen::VectorXd currentU = u0;
    
    while (t < m_t1) {
      // 如果下一步超出终止时间，调整步长
      if (t + step > m_t1) {
        step = m_t1 - t;
      }
      
      bool stepAccepted = false;
      
      while (!stepAccepted) {
        // (a) 使用两种RK方法进行推进
        Y.setZero();
        for (int j = 0; j < m_s; ++j) {
          Eigen::VectorXd tmp = currentU;
          for (int k = 0; k < j; ++k) {
            tmp += step * m_A(j, k) * Y.row(k);
          }
          Y.row(j) = f(tmp, t + m_c(j) * step);
        }
        
        // 计算两种方法的结果
        Eigen::VectorXd u = currentU;
        Eigen::VectorXd uh = currentU;
        
        for (int j = 0; j < m_s; ++j) {
          u += step * m_bl(j) * Y.row(j);   // 低阶方法
          uh += step * m_bh(j) * Y.row(j);  // 高阶方法
        }
        
        // (b) 计算误差指标
        double Eind = 0.0;
        for (int i = 0; i < m_dim; ++i) {
          // 计算每个分量的容忍度
          double eps_i = Eabs + std::abs(currentU(i)) * Erel;
          // 计算归一化误差
          double norm_error = std::abs(uh(i) - u(i)) / eps_i;
          Eind += norm_error * norm_error;
        }
        Eind = std::sqrt(Eind / m_dim);
        
        // (c) 计算新的步长
        double newStep = step * std::min(rhoMax, 
                               std::max(rhoMin, 
                                       rho * std::pow(1.0/Eind, 1.0/(q+1.0))));
        
        // (d) 判断是否接受这一步
        if (Eind <= 1.0) {
          // 接受这一步
          t += step;
          currentU = uh; // 使用高精度解
          
          // 保存结果
          solution.push_back(currentU);
          timePoints.push_back(t);
          
          step = newStep;
          // std::cout << "Step accepted: t = " << t << ", step = " << step << std::endl;
          stepAccepted = true;
        } else {
          // 拒绝这一步，使用较小的步长重试
          step = newStep;
        }
      }
    }
    
    // 将结果转移到m_u矩阵中
    m_u.resize(solution.size(), m_dim);
    for (size_t i = 0; i < solution.size(); ++i) {
      m_u.row(i) = solution[i];
    }
  }

  // Getters for time parameters
  double getDt() const { return m_dt; }
  double getT0() const { return m_t0; }
  double getT1() const { return m_t1; }
  int getS() const { return m_s; }

protected:
  int m_s;  // s-step
  int m_pl; // pl-precision 低阶方法的阶数
  int m_ph; // ph-precision 高阶方法的阶数
  
  // RK方法系数
  Eigen::MatrixXd m_A;
  Eigen::VectorXd m_c;
  Eigen::VectorXd m_bl; // 低阶方法权重
  Eigen::VectorXd m_bh; // 高阶方法权重
};

#endif